from langchain.chains.retrieval_qa.base import RetrievalQA
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain




def get_chat_chain_memory(llm, vector_store):
    # ② 存储历史记录
    # 参考官网链接：https://python.langchain.com/docs/use_cases/question_answering/how_to/chat_vector_db
    # 用于缓存或者保存对话历史记录的对象
    memory = ConversationBufferMemory(
        memory_key='chat_history', return_messages=True)


    # ③ 对话链
    conversation_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=vector_store.as_retriever(),
        memory=memory
    )
    # self.conversation_chain = conversation_chain
    return conversation_chain

def get_chat_chain_RetrievalQA(llm, vector_store):
    # 创建RetrievalQA实例
    conversation_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store.as_retriever(),
        return_source_documents=True,
        verbose=True,
    )
    return conversation_chain

